import pygame
import sys
import math

# 初始化 Pygame
pygame.init()

# 设置窗口尺寸
WIDTH, HEIGHT = 1000, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("手柄控制可视化")

# 初始化字体
font = pygame.font.SysFont('SimHei', 18)
title_font = pygame.font.SysFont('SimHei', 24, bold=True)

# 初始化手柄
pygame.joystick.init()
joysticks = [pygame.joystick.Joystick(i) for i in range(pygame.joystick.get_count())]
for joystick in joysticks:
    joystick.init()

# 颜色定义
BACKGROUND = (25, 25, 35)
TEXT_COLOR = (220, 220, 220)
HIGHLIGHT = (70, 130, 180)
AXIS_ACTIVE = (70, 180, 130)
AXIS_INACTIVE = (50, 80, 60)
BUTTON_ACTIVE = (220, 80, 80)
BUTTON_INACTIVE = (80, 50, 50)
HAT_COLOR = (180, 180, 70)
STICK_COLOR = (100, 150, 200)

def draw_joystick_visualization(joystick, x, y, width, height):
    """绘制手柄可视化界面"""
    # 绘制背景
    pygame.draw.rect(screen, (40, 40, 50), (x, y, width, height), border_radius=15)
    pygame.draw.rect(screen, (60, 60, 70), (x, y, width, height), 2, border_radius=15)
    
    center_x = x + width // 2
    center_y = y + height // 2
    
    # 绘制左摇杆
    left_stick_x = joystick.get_axis(0) if joystick.get_numaxes() > 0 else 0
    left_stick_y = joystick.get_axis(1) if joystick.get_numaxes() > 1 else 0
    draw_stick(center_x - 150, center_y - 30, left_stick_x, left_stick_y, "左摇杆")
    
    # 绘制右摇杆
    right_stick_x = joystick.get_axis(2) if joystick.get_numaxes() > 2 else 0  # 修正为轴2
    right_stick_y = joystick.get_axis(3) if joystick.get_numaxes() > 3 else 0  # 修正为轴3
    draw_stick(center_x + 150, center_y - 30, right_stick_x, right_stick_y, "右摇杆")
    
    # 绘制方向键
    if joystick.get_numhats() > 0:
        hat_value = joystick.get_hat(0)
        draw_dpad(center_x - 150, center_y + 120, hat_value)
    
    # 绘制主要按钮
    draw_buttons(center_x + 150, center_y + 120, joystick)
    
    # 绘制肩键和扳机键
    draw_triggers(center_x, y + 30, joystick)

def draw_stick(x, y, stick_x, stick_y, label):
    """绘制摇杆"""
    # 绘制摇杆底座
    pygame.draw.circle(screen, (60, 60, 70), (x, y), 40)
    pygame.draw.circle(screen, (80, 80, 90), (x, y), 40, 2)
    
    # 绘制摇杆位置
    stick_pos_x = x + stick_x * 30
    stick_pos_y = y + stick_y * 30
    pygame.draw.circle(screen, STICK_COLOR, (stick_pos_x, stick_pos_y), 15)
    
    # 绘制中心到当前位置的连线
    pygame.draw.line(screen, STICK_COLOR, (x, y), (stick_pos_x, stick_pos_y), 2)
    
    # 绘制标签
    text = font.render(f"{label}: ({stick_x:.2f}, {stick_y:.2f})", True, TEXT_COLOR)
    screen.blit(text, (x - text.get_width() // 2, y + 60))

def draw_dpad(x, y, hat_value):
    """绘制方向键 - 修复上下相反问题"""
    # 绘制方向键背景
    pygame.draw.rect(screen, (60, 60, 70), (x - 30, y - 30, 60, 60), border_radius=5)
    
    # 绘制方向键中心
    pygame.draw.circle(screen, (80, 80, 90), (x, y), 10)
    
    # 绘制按下的方向 - 修复y轴方向
    dx, dy = hat_value
    if dx != 0 or dy != 0:
        # 注意：Pygame中hat的y轴是反的，所以这里要取反
        arrow_x = x + dx * 20
        arrow_y = y - dy * 20  # 这里取反修复上下问题
        pygame.draw.line(screen, HAT_COLOR, (x, y), (arrow_x, arrow_y), 3)
        pygame.draw.circle(screen, HAT_COLOR, (arrow_x, arrow_y), 5)
    
    # 绘制标签
    # text = font.render(f"方向键: {hat_value}", True, TEXT_COLOR)
    # screen.blit(text, (x - text.get_width() // 2, y + 50))

def draw_buttons(x, y, joystick):
    """绘制主要按钮 - 修正按钮布局"""
    # ABXY 按钮布局（修正为：上Y，下A，左X，右B）
    buttons = [
        # (相对位置x, 相对位置y, 按钮索引, 标签, 颜色)
        (0, 40, 0, "A", (80, 180, 80)),     # 下 - 绿色A
        (40, 0, 1, "B", (180, 80, 80)),     # 右 - 红色B
        (0, -40, 2, "Y", (80, 80, 180)),    # 上 - 蓝色Y
        (-40, 0, 3, "X", (180, 180, 80)),   # 左 - 黄色X
    ]
    
    for dx, dy, btn_idx, label, color in buttons:
        btn_x = x + dx
        btn_y = y + dy
        
        # 检查按钮是否按下
        is_pressed = joystick.get_button(btn_idx) if joystick.get_numbuttons() > btn_idx else False
        
        # 绘制按钮
        btn_color = color if is_pressed else (color[0]//2, color[1]//2, color[2]//2)
        pygame.draw.circle(screen, btn_color, (btn_x, btn_y), 20)
        pygame.draw.circle(screen, (200, 200, 200), (btn_x, btn_y), 20, 2)
        
        # 绘制按钮标签
        text = font.render(label, True, (255, 255, 255))
        screen.blit(text, (btn_x - text.get_width() // 2, btn_y - text.get_height() // 2))

def draw_triggers(center_x, y, joystick):
    """绘制肩键和扳机键 - 修复文字重叠"""
    # 左扳机
    left_trigger = joystick.get_axis(4) if joystick.get_numaxes() > 4 else 0  # 修正轴索引
    if left_trigger < 0:  # 有些手柄扳机键范围是 -1 到 1
        left_trigger = (left_trigger + 1) / 2
    draw_trigger(center_x - 120, y, left_trigger, "左扳机", (180, 100, 80))
    
    # 右扳机
    right_trigger = joystick.get_axis(5) if joystick.get_numaxes() > 5 else 0  # 修正轴索引
    if right_trigger < 0:
        right_trigger = (right_trigger + 1) / 2
    draw_trigger(center_x + 120, y, right_trigger, "右扳机", (80, 100, 180))  # 调整位置避免重叠
    
    # 肩键
    left_bumper = joystick.get_button(4) if joystick.get_numbuttons() > 4 else False
    right_bumper = joystick.get_button(5) if joystick.get_numbuttons() > 5 else False
    draw_bumper(center_x - 120, y + 80, left_bumper, "左肩键", (180, 100, 80))  # 调整位置
    draw_bumper(center_x + 120, y + 80, right_bumper, "右肩键", (80, 100, 180))  # 调整位置

def draw_trigger(x, y, value, label, color):
    """绘制扳机键"""
    height = 80
    width = 30
    
    # 绘制背景
    pygame.draw.rect(screen, (60, 60, 70), (x - width//2, y, width, height), border_radius=5)
    
    # 绘制激活部分
    active_height = int(value * height)
    pygame.draw.rect(screen, color, (x - width//2, y + height - active_height, width, active_height), border_radius=5)
    
    # 绘制边框
    pygame.draw.rect(screen, (200, 200, 200), (x - width//2, y, width, height), 2, border_radius=5)
    
    # 绘制标签
    text = font.render(f"{label}: {value:.2f}", True, TEXT_COLOR)
    screen.blit(text, (x - text.get_width() // 2, y + height + 30))

def draw_bumper(x, y, is_pressed, label, color):
    """绘制肩键"""
    width = 60
    height = 20
    
    # 绘制肩键
    btn_color = color if is_pressed else (color[0]//2, color[1]//2, color[2]//2)
    pygame.draw.rect(screen, btn_color, (x - width//2, y, width, height), border_radius=5)
    pygame.draw.rect(screen, (200, 200, 200), (x - width//2, y, width, height), 2, border_radius=5)
    
    # 绘制标签
    text = font.render(f"{label}: {'按下' if is_pressed else '释放'}", True, TEXT_COLOR)
    screen.blit(text, (x - text.get_width() // 2, y + height + 30))

def main():
    clock = pygame.time.Clock()
    
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    sys.exit()
        
        # 更新手柄状态
        for joystick in joysticks:
            pygame.event.pump()
        
        # 绘制界面
        screen.fill(BACKGROUND)
        
        # 显示连接状态
        if not joysticks:
            text = title_font.render("未检测到手柄，请连接手柄", True, TEXT_COLOR)
            screen.blit(text, (WIDTH//2 - text.get_width()//2, HEIGHT//2 - text.get_height()//2))
        else:
            # 绘制每个手柄的可视化
            for i, joystick in enumerate(joysticks):
                title = title_font.render(f"手柄 {i}: {joystick.get_name()}", True, HIGHLIGHT)
                screen.blit(title, (WIDTH//2 - title.get_width()//2, 20))
                draw_joystick_visualization(joystick, 50, 70, WIDTH - 100, HEIGHT - 120)
        
        # 显示说明
        help_text = font.render("按 ESC 退出程序", True, TEXT_COLOR)
        screen.blit(help_text, (WIDTH - help_text.get_width() - 20, HEIGHT - 30))
        
        pygame.display.flip()
        clock.tick(60)

if __name__ == "__main__":
    main()